package main

import "fmt"

// 一种自平衡二叉搜索树，任何节点的两个子树的高度差不超过1。
/**
A右旋:A节点的左节点B的右子树(如果有)变成A节点的左子树;A节点变成B节点的右节点;B节点变成A节点;
A左旋:A节点的右节点C的左子树(如果有)变成A节点的右子树;A节点变成C节点的左子树;C节点变成A节点;
LL 型  插入的节点在节点A的左节点的左子树上(树高(左子树高-右子树高)大于1),这时候需要将失衡节点A右旋
RR 型: 插入的节点在节点A的右节点的右子树上(树高(左子树-右子树高度)小于-1),这时候需要将失衡节点A左旋
LR 型: 插入的节点在节点A的左节点的右子树上(树高(左子树高-右子树高)大于1),这时候先将失衡节点的左子节点左旋 ,然后再将失衡节点右旋
RL 型: 插入的节点在节点A的右节点的左子树上(树高(左子树-右子树高度)小于-1),这时候需要将A节点的右节点右旋,然后再将失衡节点左旋
*/
func main() {
	var root *TreeNode
	values := []int{10, 20, 30, 40, 50, 25}
	for _, value := range values {
		root = root.insert(value)
	}

	fmt.Println("in-order traversal of the AVL tree:")
	root.inOrderTraversal()
}

type TreeNode struct {
	Left   *TreeNode
	Value  int
	Right  *TreeNode
	Height int
}

func (n *TreeNode) getTreeHeight() int {
	if n != nil {
		return n.Height
	}
	return 0
}

func (n *TreeNode) updateTreeHeight() {
	leftHeight := n.Left.getTreeHeight()
	rightHeight := n.Right.getTreeHeight()
	if leftHeight > rightHeight {
		n.Height = leftHeight + 1
	} else {
		n.Height = rightHeight + 1
	}
}

// 右旋操作
func (n *TreeNode) rotateRight() *TreeNode {
	left := n.Left
	n.Left = left.Right
	left.Right = n
	left.updateTreeHeight()
	n.updateTreeHeight()
	return left
}

// 左旋
func (n *TreeNode) rotateLeft() *TreeNode {
	right := n.Right
	n.Right = right.Left
	right.Left = n
	n.updateTreeHeight()
	right.updateTreeHeight()
	return right
}

// 平衡
func (n *TreeNode) balance() *TreeNode {
	n.updateTreeHeight()
	balanceFactor := n.Left.getTreeHeight() - n.Right.getTreeHeight()
	if balanceFactor > 1 {
		if n.Left.Left.getTreeHeight() < n.Left.Right.getTreeHeight() {
			// LR
			n.Left = n.Left.rotateLeft()
		}
		return n.rotateRight()
	}
	if balanceFactor < -1 {
		if n.Right.Right.getTreeHeight() < n.Right.Left.getTreeHeight() {
			// RL
			n.Right = n.Right.rotateRight()
		}
		return n.rotateLeft()
	}
	return n
}

func (n *TreeNode) insert(value int) *TreeNode {
	if n == nil {
		return &TreeNode{Value: value, Height: 1}
	}
	if value < n.Value {
		n.Left = n.Left.insert(value)
	} else {
		n.Right = n.Right.insert(value)
	}
	return n.balance()
}

func (n *TreeNode) inOrderTraversal() {
	if n != nil {
		n.Left.inOrderTraversal()
		fmt.Println(n.Value)
		n.Right.inOrderTraversal()
	}
}
